# LeetCode 310. 最小高度树
# https://leetcode.cn/problems/minimum-height-trees/description
from collections import deque
from math import inf
from typing import List


# BFS 遍历每一个节点 超时 用例通过: 65/71
class Solution:
    def findMinHeightTrees(self, n: int, edges: List[List[int]]) -> List[int]:
        if n == 1:
            return [0]

        adj = [[] for i in range(n)]
        for edge in edges:
            node_0, node_1 = edge
            adj[node_0].append(node_1)
            adj[node_1].append(node_0)

        def bfs(s):
            d = [inf for i in range(n)]
            d[s] = 0
            q = deque()
            q.append(s)
            while q:
                u = q.popleft()
                for v in adj[u]:
                    if d[v] == inf:
                        d[v] = d[u] + 1
                        q.append(v)
            return d

        min_height = inf
        result = []
        for root in range(n):
            d = bfs(root)
            if min_height == max(d):
                result.append(root)
            elif min_height > max(d):
                min_height = min(min_height, max(d))
                result = [root]
        return result


# BFS 遍历叶子节点
class Solution2:
    def findMinHeightTrees(self, n: int, edges: List[List[int]]) -> List[int]:
        if n == 1:
            return [0]
        elif n == 2:
            return [0, 1]

        adj = [[] for i in range(n)]
        for edge in edges:
            node_0, node_1 = edge
            adj[node_0].append(node_1)
            adj[node_1].append(node_0)

        leaves = [i for i in range(n) if len(adj[i]) == 1]

        while True:
            leaves_tmp = []
            for u in leaves:
                neighbours = {v for v in adj[u] if v not in leaves}
                for neighbour in neighbours:
                    adj[u].remove(neighbour)
                    adj[neighbour].remove(u)
                    if len(adj[neighbour]) == 1:
                        leaves_tmp.append(neighbour)
            if not leaves_tmp:
                leaves = [i for i in range(n) if len(adj[i]) == 1]
            else:
                leaves = leaves_tmp

            if (len(leaves) == 2 and adj[leaves[0]][0] == leaves[1] and adj[leaves[1]][0] == leaves[0]) \
                    or len(leaves) == 1:
                return list(leaves)


if __name__ == '__main__':
    n = 4
    edges = [[1, 0], [1, 2], [1, 3]]
    n = 6
    edges = [[3, 0], [3, 1], [3, 2], [3, 4], [5, 4]]
    # n = 1
    # edges = []
    # n = 2
    # edges = [[0, 1]]
    # n = 6
    # edges = [[0,1],[0,2],[0,3],[3,4],[4,5]]
    print(Solution().findMinHeightTrees(n, edges))
    print(Solution2().findMinHeightTrees(n, edges))
